#!/home/zhonc0a/Enthought/Canopy_64bit/User/bin/python
import matplotlib
matplotlib.use('Agg')
import pylab
import re
import argparse
import sys
from scipy.ndimage.filters import gaussian_filter1d
class dossComb():
    def __init__(self,title):
        self.eneRange=''
        self.dpi=150
        self.title=''.join(re.split('(\.)',title)[:-2])
        self.outFileName=title
        self.dossRange=''
        self.gaussian=None
        self.combinedDoss=[]
         
    def readDOSS(self,dossFileName):
        self.dossData=[]
        dossRawData=open(dossFileName)
        for line in dossRawData:
            if re.search('-?\d+E[-+]\d+',line):
                self.dossData.append(map(float,line.split()))
        self.npdos=len(self.dossData[0])-1
        self.dossX=pylab.array(map(float,pylab.array(self.dossData).transpose()[0] ))*27.211396
        self.dossY=pylab.array(map(float,pylab.array(self.dossData).transpose()[-1] ))/27.211396
        if self.gaussian>0:
            self.FWHM=((self.dossX[0]-self.dossX[-1])/len(self.dossX)*self.gaussian)*2.355
            WdossY=gaussian_filter1d(self.dossY,self.gaussian,truncate=100.0,mode='constant')
            self.dossY=WdossY
        self.dossYmin=self.dossY.min()
        self.dossYmax=self.dossY.max()
        self.dossXmin=self.dossX.min()
        self.dossXmax=self.dossX.max()
    
    def combineDOSS(self,multi):
        if self.combinedDoss==[]:
            self.combinedDoss=self.dossData
            self.combinedY=self.dossY*multi
        else:
            for i in range(len(self.combinedDoss)):
                if self.combinedDoss[i][0]*27.211396 == self.dossX[i]:
                    self.combinedY[i] = self.combinedY[i]+self.dossY[i]*multi
                else:
                    print "the energy of the doss file are not equal"+str(self.combinedDoss[i][0])+" "+str(self.dossX[i])+", can not combine"
                    sys.exit()
            self.dossY=self.combinedY
            self.dossYmin=self.dossY.min()
            self.dossYmax=self.dossY.max()
        
        
    def plotDoss(self,style,name,multi=1):
        if self.dossRange:
            self.dossXmin=float(self.dossRange.split(',')[0])*-1
            self.dossXmax=float(self.dossRange.split(',')[1])
        print self.dossY
        print self.dossX
        pylab.plot(self.dossY*multi,self.dossX,style)

        pylab.axis([self.dossYmin,self.dossYmax,self.dossXmin, self.dossXmax])
        #write doss date to file
        dossDATA=open('DOSSDATA','w')
        for i in range(0,len(self.dossX)):
            dossDATA.write('%10f %f\n' % (self.dossX[i],self.dossY[i]*4))
        dossDATA.close()
        pylab.ylabel('Energy (eV)')
        pylab.xlabel('States/eV/Cell')
        pylab.savefig(self.title+'_doss_'+name+'.png',dpi=self.dpi,bbox_inches='tight',format='png')
        pylab.savefig(self.title+'_doss_'+name+'.svg',bbox_inches='tight',format='svg')
        
parser=argparse.ArgumentParser(description='combine two doss file and plot')
parser.add_argument("-r","--range",help='range for draw band structure, e.g. 11,2 means draw form -11 eV to 2 eV, default:all the point',type=str)
parser.add_argument("-g","--gaussian",help='gaussian widen for the DOS, set the sigma, default:None',type=int )
parser.add_argument("-m","--multi",help='multiplier for each doss, default:all 1',type=str)
#parser.add_argument("-p","--pdos",help='input three tuple. The first one is the column to plot, the second one is the Label, and the third one is the style, e.g. 1,2,3;S,Px,Py;r,g,b, default:None',type=str )
parser.add_argument("dossFile",nargs=2)
args=parser.parse_args()

if args.multi:
    multi=map(int,args.multi.split(','))
else:
    multi=[1]*len(args.dossFile)   
print multi  
dossgen=dossComb(args.dossFile[0])
if args.range:
    dossgen.dossRange=args.range
dossgen.gaussian=args.gaussian
dossgen.readDOSS(args.dossFile[0])
dossgen.plotDoss('r-','doss0',multi=4)
dossgen.combineDOSS(multi[0])
dossgen.readDOSS(args.dossFile[1])
dossgen.plotDoss('g-','doss1')
dossgen.combineDOSS(multi[1])

dossgen.plotDoss('b-','combined')
if dossgen.gaussian>0:
    print dossgen.title+': FWHM is '+str(dossgen.FWHM)+'(eV)'
    #print bandgen.HOMOband
    
    
    #if last_max==pylab.array(map(float,line)).min()
    #last_max=pylab.array(map(float,line)).max()
    #last_min=pylab.array(map(float,line)).min()
#    





